#include <iostream>
#include <chrono>
#include <vector>
#include <math.h>

#include <ros/ros.h>
#include <message_filters/subscriber.h>
#include <message_filters/time_synchronizer.h>
#include <visualization_msgs/Marker.h>
#include <visualization_msgs/MarkerArray.h>

#include <std_msgs/Int8.h>
#include <nav_msgs/Odometry.h>
#include <nav_msgs/OccupancyGrid.h>

// PCL specific includes
#include <pcl/io/pcd_io.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/common/transforms.h>
#include <pcl_ros/transforms.h>
#include "pcl_ros/impl/transforms.hpp"

// path
#include <nav_msgs/Path.h>

// opencv
#include <opencv2/opencv.hpp>

#include <object_tracking/trackbox.h>
#include <object_tracking/resbox.h>
#include <object_tracking/resboxes.h>

#include "ImmUkf.hpp"
#include "vis_bbox.h"

using namespace std;
using namespace Eigen;
using namespace pcl;


string lidar_name = "velo_link";
string global_name = "global";
string params_path;
bool use_costmap = false;

ros::Publisher vis_pub;
ros::Publisher text_pub;
ros::Publisher costmap_pub;
ros::Publisher ego_pub;
ros::Publisher path_pub;

struct DyncAvoiding{

	int8_t STAY = 1;
	int8_t AVOID = 2;
	
	int8_t mode = STAY;
	int no_obj = 0;
	int release_cnt = 20;
	ros::Publisher avoiding_cmd_pub;
	ros::Publisher switch_mode_pub;
	
	void incNoObj(){
		no_obj++;
		if(no_obj > release_cnt){
			no_obj = 0;
			switch_mode();
		}
	}

	void switch_mode(){
		std_msgs::Int8 msg;
		if(mode == STAY){
			mode = AVOID;
			msg.data = mode;
			switch_mode_pub.publish(msg);
		}else{
			mode = STAY;
			msg.data = mode;
			switch_mode_pub.publish(msg);
		}
	}
	
	double car_radius = 1.0;
	double max_vel = 8.0;
	double max_ang = M_PI / 3;
	double ang_mul = 3.0;
	double vel_mul = 3.0;
	void genAvoidTwist(const vector<vector<Vector3d>>& preTraj, double colliProb)
	{
		size_t num = preTraj.size();
		if(num == 0)	return;

		Vector2d combineVec;
		combineVec.fill(0.0);

		for (size_t i = 0; i < num; i++)
		{
			int firstInCircleId = -1;
			int lastInCircleId = -1;
			// find a section of traj that in circle
			for (size_t j = 0; j < preTraj[i].size(); j++)
			{
				if(preTraj[i][j].head<2>().norm() < car_radius)
				{
					firstInCircleId = j;
					break;
				}
			}

			if(firstInCircleId >= 0)
			{
				int firstOutOfCircleId = -1;
				for (size_t j = firstInCircleId+1; j < preTraj[i].size(); j++)
				{
					if(preTraj[i][j].head<2>().norm() > car_radius)
					{
						firstOutOfCircleId = j;
						break;
					}
				}

				if(firstOutOfCircleId == -1)	firstOutOfCircleId = preTraj[i].size() - 1;

				Vector2d tmp_front = preTraj[i][firstOutOfCircleId].head<2>();
				Vector2d tmp_back = preTraj[i][firstOutOfCircleId-1].head<2>();
				tmp_back = tmp_front - tmp_back;
				tmp_back /= tmp_back.norm();
				combineVec += tmp_back;
			}
			
		}

		// max vel assume 5.0
		double vel = (combineVec(0) > 0 ? 1.0 : -1.0) * std::min(colliProb * max_vel * vel_mul, max_vel);
		double angle = std::atan2(combineVec(1), combineVec(0));
		angle = (std::cos(angle+M_PI_2) > std::cos(angle-M_PI_2) ? angle+M_PI_2 : angle-M_PI_2);
		
		while (angle > 2*M_PI)	angle -= 2 * M_PI;
		while (angle < -2*M_PI) angle += 2 * M_PI;

		geometry_msgs::Twist v;
		v.linear.x = vel;
		v.linear.y = 0.0;
		v.linear.z = 0.0;
		
		v.angular.x = 0.0;
		v.angular.y = 0.0;
		double yaw_sign = (angle > 0 ? 1.0 : -1.0);
		angle = std::min(std::abs(colliProb * angle * ang_mul), max_ang);
		v.angular.z = yaw_sign * angle;
		
		avoiding_cmd_pub.publish(v);
		ROS_WARN("avoiding (%f, %f)", vel, angle);
	}

}dyav;


// some parameters of cost map
nav_msgs::OccupancyGrid og;
double resolution = 0.1; 	// 0.1m/pixel
int map_size = 400; 		// local map 400x400
double occ_color[2] = {20, 100};

// rearrange the input
vector<PointCloud<PointXYZ>> formatTrackbox(const object_tracking::trackbox& input)
{
	
	vector<PointCloud<PointXYZ>> bBoxes;
	int box_num = input.box_num;
	bBoxes.reserve(box_num);

	for(int box_i = 0;box_i < box_num; box_i++)
	{
		PointCloud<PointXYZ> oneBbox;
		PointXYZ o;
		o.x = input.x1[3*box_i];
		o.y = input.x1[3*box_i + 1];
		o.z = input.x1[3*box_i + 2];
		oneBbox.push_back(o);
		o.x = input.x2[3*box_i];
		o.y = input.x2[3*box_i + 1];
		o.z = input.x2[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.x3[3*box_i + 0];
		o.y = input.x3[3*box_i + 1];
		o.z = input.x3[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.x4[3*box_i + 0];
		o.y = input.x4[3*box_i + 1];
		o.z = input.x4[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.y1[3*box_i + 0];
		o.y = input.y1[3*box_i + 1];
		o.z = input.y1[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.y2[3*box_i + 0];
		o.y = input.y2[3*box_i + 1];
		o.z = input.y2[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.y3[3*box_i + 0];
		o.y = input.y3[3*box_i + 1];
		o.z = input.y3[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.y4[3*box_i + 0];
		o.y = input.y4[3*box_i + 1];
		o.z = input.y4[3*box_i + 2];   
		oneBbox.push_back(o);
		bBoxes.emplace_back(oneBbox);
	}

	return bBoxes;
}

void formatResboxes(const object_tracking::resboxes& input, vector<int>& id, vector<Vector3d>& color, vector<PointCloud<PointXYZ>>& trackingBBoxes,  vector<vector<Vector3d>>& predTraj, vector<Vector3d>& egoPredTraj, vector<Vector3d>& xyz)
{
	int n = input.rb.size();
	trackingBBoxes.reserve(n);
	predTraj.reserve(n);
	id.reserve(n);
	color.reserve(n);
	xyz.reserve(n);
	for(auto&& e : input.rb)
	{
		id.emplace_back(e.id);
		Vector3d tmp_color, tmp_xyz;
		tmp_color << e.color[0], e.color[1], e.color[2];
		color.emplace_back(std::move(tmp_color));
		tmp_xyz << e.xyz[0], e.xyz[1], e.xyz[2];
		xyz.emplace_back(std::move(tmp_xyz));

		PointCloud<PointXYZ> pc;
		pc.reserve(8);
		for(int i=0; i<8; i++)
		{
			PointXYZ o;
			o.x = e.bx[i];
			o.y = e.by[i];
			o.z = e.bz[i];
			pc.push_back(std::move(o));
		}
		trackingBBoxes.emplace_back(std::move(pc));

		vector<Vector3d> tr;
		tr.reserve(e.traj_x.size());
		for(int i=0; i<e.traj_x.size(); i++)
		{
			Vector3d tmp;
			tmp << e.traj_x[i], e.traj_y[i], e.traj_z[i];
			tr.emplace_back(std::move(tmp));
		}
		predTraj.emplace_back(std::move(tr));
	}

	const auto& et = input.ego;
	for(int i=0; i<et.traj_x.size(); i++){
		Vector3d p;
		p << et.traj_x[i], et.traj_y[i], et.traj_z[i];
		egoPredTraj.emplace_back(p);
	}
}

void createLocalCostMap(const vector<Vector3d>& xyz, const vector<vector<Vector3d>>& traj, vector<int>& cost_map, double& colli_prob)
{
	double ofs = map_size / 2;

	cv::Mat cost_map_cv(map_size, map_size, CV_8UC1, cv::Scalar(0));

	for(int i=0; i<xyz.size(); i++)
	{
		//double r_x = xyz[i](0) / 2;
		//double r_y = xyz[i](1) / 2;
		//int r = sqrt(r_x * r_x + r_y * r_y) / resolution;
		double r = 1.5 * xyz[i](1) / 2 / resolution;
		
		const auto& etj = traj[i];
		Vector2d lastp;
		int tjSize = etj.size();
		for(int j=0; j<tjSize; j++)
		{
			int color = occ_color[0] + 1.0 * (tjSize - j) / tjSize * (occ_color[1] - occ_color[0]);
			if(j == 0){
				lastp << etj[j](0)/resolution + ofs, etj[j](1)/resolution + ofs;
				cv::circle(cost_map_cv, cv::Point(lastp(0), lastp(1)), r, color, cv::FILLED);
				continue;	
			}

			// draw rect
			Vector2d curp;
			curp << etj[j](0)/resolution + ofs, etj[j](1)/resolution + ofs;
			
			Vector2d dir = curp - lastp;
			double len = dir.norm() / resolution;
			double yaw = atan2(dir(1), dir(0));
			Matrix<double, 2, 4> vert;
			vert.col(0) = Vector2d(0, r);
			vert.col(1) = Vector2d(0, -r);
			vert.col(2) = Vector2d(dir.norm(), -r);
			vert.col(3) = Vector2d(dir.norm(), r);
			//vert.col(0) = Vector2d(r, 0);
			//vert.col(1) = Vector2d(-r, 0);
			//vert.col(2) = Vector2d(r, dir.norm());
			//vert.col(3) = Vector2d(-r, dir.norm());
			
			Matrix2d rot;
			rot << cos(yaw), -sin(yaw), sin(yaw), cos(yaw);
			vert = rot * vert;
			
			vector<cv::Point> cv_vert;
			for(int k=0; k<4; k++){
				cv::Point v(lastp(0)+vert(0, k), lastp(1)+vert(1, k));
				cv_vert.emplace_back(v);
			}
			//std::cout << vert << endl;

			cv::fillConvexPoly(cost_map_cv, cv_vert, color);
			cv::circle(cost_map_cv, cv::Point(curp(0), curp(1)), r, color, cv::FILLED);
			lastp = curp;
			
		}
	}

	// make cv mat to vector
	for(int i=0; i<map_size; i++){
		for(int j=0; j<map_size; j++){
			cost_map[i * map_size + j] = cost_map_cv.at<uchar>(i, j);
		}
	}
	
	// check collision
	int ego_r = 0.5 / resolution;
	int accumulate_occ = 0;
	int cnt = 0;
	for(int cx=-ego_r; cx < ego_r; cx++){
		for(int cy=-ego_r; cy < ego_r; cy++){
			accumulate_occ += cost_map_cv.at<uchar>(ofs + cx, ofs + cy);
			cnt++;
		}
	}
	colli_prob = 1.0 * accumulate_occ / cnt / occ_color[1];
			
}

void detectBoxesVis(const object_tracking::trackbox &cbox)
{
	const ros::Time& input_time = cbox.header.stamp;
	vector<PointCloud<PointXYZ>> cbox_vec = formatTrackbox(cbox);
	
	// pub cluster box
	visualization_msgs::Marker line_list; //将候选框8个点连线
	vector<double> color = {0, 1, 0};
	//line_list = vis_utils::getBBox(0, input_time, lidar_name, color, cbox_vec); 
	line_list = vis_utils::getBBox(0, input_time, global_name, color, cbox_vec); 
	vis_pub.publish(line_list);
}

void resultsVis(const object_tracking::resboxes& tbox)
{
	const ros::Time& input_time = tbox.header.stamp;
	
	vector<PointCloud<PointXYZ>> trackingBBoxes;
	vector<vector<Vector3d>> predTraj;
	vector<int> box_id;
	vector<Vector3d> box_color;
	vector<Vector3d> egoPredTraj;
	vector<Vector3d> box_xyz;
	formatResboxes(tbox, box_id, box_color, trackingBBoxes, predTraj, egoPredTraj, box_xyz);

	visualization_msgs::Marker tBBoxesMarker; //将候选框8个点连线
	vector<PointCloud<PointXYZ>> tmpTrBox(1);

	// pub trajectory
	vis_utils::rosTraj rT(lidar_name, input_time);
	for(int i=0; i<trackingBBoxes.size(); i++)
	{
		vector<double> color(3, 0);
		for(int j=0; j<3; j++) 	color[j] = box_color[i](j);

		int id = box_id[i] + 1;
		tmpTrBox[0] = trackingBBoxes[i];
		tBBoxesMarker = vis_utils::getBBox(id, input_time, lidar_name, color, tmpTrBox); 
		vis_utils::make_visualize_traj(id, input_time, predTraj[i], rT);
		
		vis_pub.publish(tBBoxesMarker);
		vis_pub.publish(rT.line);	
	}
	// pub ego
	if(egoPredTraj.size() > 1){
		rT.line.header.frame_id = global_name;
		rT.line.color.g = 0;
		vis_utils::make_visualize_traj(-1, input_time, egoPredTraj, rT);
		vis_pub.publish(rT.line);
	}

	try{
		visualization_msgs::MarkerArray text_marker = vis_utils::getLabels(tbox);
		text_pub.publish(text_marker);
	}catch(exception& e){
		// ROS_WARN("visualization error : %s", e.what());
	}

	if(use_costmap){
		vector<int> cost_map(map_size*map_size, 0);
		double colli_prob = 0;
		createLocalCostMap(box_xyz, predTraj, cost_map, colli_prob);
		// pub cost map
		og.data.insert(og.data.end(), cost_map.begin(), cost_map.end());
		costmap_pub.publish(og);
		og.data.clear();

		char colli_prob_str[32];
		sprintf(colli_prob_str, "%.2f", colli_prob);
		string msg = "Collision Prob: " + string(colli_prob_str);
		visualization_msgs::Marker alarm = vis_utils::getText(lidar_name, msg, {1.0, 0, 0});
		ego_pub.publish(alarm);

		if(dyav.mode == dyav.AVOID){
			if(colli_prob < 1e-2){
				dyav.incNoObj();
			}
			dyav.genAvoidTwist(predTraj, colli_prob);
			
		}else{
			if(colli_prob >= 1e-2){
				dyav.switch_mode();
				dyav.genAvoidTwist(predTraj, colli_prob);
			}
		}
	}

}

//void res_cb(const object_tracking::trackbox::ConstPtr& cbox, const object_tracking::resboxes::ConstPtr& tbox)
//{
	
	//const ros::Time& input_time = cbox->header.stamp;

	//vector<PointCloud<PointXYZ>> cbox_vec = formatTrackbox(*cbox);

	//vector<PointCloud<PointXYZ>> trackingBBoxes;
	//vector<vector<Vector3d>> predTraj;
	//vector<int> box_id;
	//vector<Vector3d> box_color;
	//formatResboxes(*tbox, box_id, box_color, trackingBBoxes, predTraj);

	//visualization_msgs::Marker tBBoxesMarker; //将候选框8个点连线
	//vector<PointCloud<PointXYZ>> tmpTrBox(1);
	
	//// pub cluster box
	//visualization_msgs::Marker line_list; //将候选框8个点连线
	//vector<double> color = {0, 1, 0};
	//line_list = vis_utils::getBBox(0, input_time, lidar_name, color, cbox_vec); 
	//vis_pub.publish(line_list);
	
	//// pub trajectory
	//vis_utils::rosTraj rT(lidar_name, input_time);
	//for(int i=0; i<trackingBBoxes.size(); i++)
	//{
		//vector<double> color(3, 0);
		//for(int j=0; j<3; j++) 	color[j] = box_color[i](j);

		//int id = box_id[i] + 1;
		//tmpTrBox[0] = trackingBBoxes[i];
		//tBBoxesMarker = vis_utils::getBBox(id, input_time, lidar_name, color, tmpTrBox); 
		//vis_utils::make_visualize_traj(id, input_time, predTraj[i], rT);
		
		//vis_pub.publish(tBBoxesMarker);
		//vis_pub.publish(rT.line);	
	//}

	//try{
		//visualization_msgs::MarkerArray text_marker = vis_utils::getLabels(*tbox);
		//text_pub.publish(text_marker);
	//}catch(exception& e){
		//ROS_WARN("visualization error : %s", e.what());
	//}

//}

void pubEgoPath(const ros::TimerEvent& event)
{
	static nav_msgs::Path ego_path;
	ego_path.header.stamp = ros::Time::now();
    ego_path.header.frame_id=global_name;

	static tf::TransformListener tran(ros::Duration(10.0));
	tf::StampedTransform ego_tf;
	tran.waitForTransform("/" + global_name, "/" + lidar_name, ros::Time(0), ros::Duration(0.1));
	try{
		tran.lookupTransform("/" + global_name, "/" + lidar_name, ros::Time(0), ego_tf);
	}catch(exception& e) 
	{
		return;
	}
	geometry_msgs::PoseStamped pose;
	tf::Vector3 tt = ego_tf.getOrigin();
	pose.pose.position.x = tt.getX();
	pose.pose.position.y = tt.getY();
	pose.pose.position.z = 5.0;
	// pose.pose.orientation.x = 0;
	// pose.pose.orientation.y = 0;
	// pose.pose.orientation.z = 0;
	// pose.pose.orientation.w = 1.0;

	pose.header.frame_id = global_name;
	pose.header.stamp = ros::Time::now();

	ego_path.poses.push_back(pose);

	path_pub.publish(ego_path);

}


int main (int argc, char** argv){
	// Initialize ROS
	ros::init (argc, argv, "vis_node");
	ros::NodeHandle nh;
	ros::NodeHandle p_nh("~");

	// get params
	nh.getParam("lidar_name", lidar_name);
	nh.getParam("global_name", global_name);
	nh.getParam("params_path", params_path);
	p_nh.getParam("use_costmap", use_costmap);

	cv::FileStorage fs(params_path, cv::FileStorage::READ);
	assert(fs.isOpened());
	fs["CostMap"]["resolution"] >> resolution;	
	fs["CostMap"]["map_size"] >> map_size;	
	fs.release();
	
	
	ros::Subscriber box_sub = nh.subscribe("track_box", 1, detectBoxesVis);
	ros::Subscriber track_sub = nh.subscribe("resbox", 1, resultsVis);
	
	ros::Timer ego_path = nh.createTimer(ros::Duration(0.5), pubEgoPath);

	//message_filters::Subscriber<object_tracking::trackbox> box_sub(nh, "track_box", 1000);
	//message_filters::Subscriber<object_tracking::resboxes> track_sub(nh, "resbox", 1000);
	//message_filters::TimeSynchronizer<object_tracking::trackbox, object_tracking::resboxes> sync(box_sub, track_sub, 1000);
	//sync.registerCallback(boost::bind(&res_cb, _1, _2));
	
	vis_pub = nh.advertise<visualization_msgs::Marker>( "vis_markers", 100);
	text_pub = nh.advertise<visualization_msgs::MarkerArray>( "text_markers", 100);
	path_pub = nh.advertise<nav_msgs::Path>("ego_path", 10);

	if(use_costmap){
		ego_pub = nh.advertise<visualization_msgs::Marker>( "ego_text_marker", 10);
		// cost map publisher 	
		costmap_pub = nh.advertise<nav_msgs::OccupancyGrid>("cost_map", 10);  
		
		dyav.avoiding_cmd_pub = nh.advertise<geometry_msgs::Twist>("/mobile_base_controller/cmd_vel", 1);  
		dyav.switch_mode_pub = nh.advertise<std_msgs::Int8>("/car_cmd_switch", 10);  

		og.header.frame_id = lidar_name;
		og.info.resolution = resolution;
		og.info.width = map_size;
		og.info.height = map_size;
		og.info.origin.position.x = (-1) * (map_size / 2.0) * resolution;
		og.info.origin.position.y = (-1) * (map_size / 2.0) * resolution;
		og.info.origin.position.z = 0;
		og.info.origin.orientation.y = 0.0;
		og.info.origin.orientation.x = 0.0;
		og.info.origin.orientation.z = 0.0;
		og.info.origin.orientation.w = 1.0;
		og.data.reserve(map_size * map_size);	
	}
	// Spin
	ros::spin ();
}
